from scapy.all import (
    Ether,
    IP,
    ARP,
    TCP,
    UDP,
    ICMP,
    sniff
)
from packet import BasicPacket
from hackrandom import *
from error import *

def check_protocol(protocol:str) -> str:
    protocol = protocol.upper()
    if protocol=='ETHRE':
        return 'Ether'
    elif protocol=='RAW':
        return 'Raw'
    elif protocol=='LOAD':
        return 'Raw'
    return protocol

class Sniffer:
    #magic method
    def __init__(self):
        self.__result__:list = []
    
    def __str__(self):
        return f'A sniffer with {len(self)} packet(s)'

    def __len__(self):
        return len(self.__result__)

    def __iter__(self):
        self.__iter_index__ = 0
        return self

    def __next__(self):
        if 0<=self.__iter_index__<len(self):
            return self.__result__[self.__iter_index__]
        raise StopIteration()

    def __getitem__(self, index:int):
        return self.__result__[index]

    def __add__(self, sniffer):
       self.__result__ += sniffer.result()
    
    def __mod__(self, protocol:str):
        output:list = []
        protocol = check_protocol(protocol.upper())

        for pakcet in self.__result__:
            try:
                pakcet[protocol]
            except:
                pass
            else:
                output.append(pakcet)
        return output
    
    def __floordiv__(self, protocol:str):
        output:list = []
        protocol = check_protocol(protocol.upper())

        for packet in self.__result__:
            try:
                output.append(packet[protocol])
            except:
                pass
        return output

    #method
    def result(self, protocol:str=None):
        if not protocol:
            return self.__result__[::]
    
    def clear(self):
        amount = len(self)
        self.__result__.clear()
        return amount

    def sniff(
            self, 
                count:int=0, 
                timeout:int=None, 
                filter:str=None, 
               **kwargs
        ):
        result = sniff(
            count=count,
            timeout=timeout,
            filter=filter,
            **kwargs
        )
        self.__result__ += result
        return result